package cn.xaut.ocr3.utils;

import com.benjaminwan.ocrlibrary.OcrResult;
import com.benjaminwan.ocrlibrary.TextBlock;
import io.github.mymonstercat.Model;
import io.github.mymonstercat.ocr.InferenceEngine;
import io.github.mymonstercat.ocr.config.ParamConfig;

import java.io.File;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

/**
 * @author dgl
 */
public class RapidOcrUtils {

    /**
     * 识别图片
     * @param file 图片文件
     * @return 识别结果
     */
    public static String recognizeImage(File file, InferenceEngine engine, ParamConfig paramConfig){
        // 识别图片
        OcrResult ocrResult = engine.runOcr(file.getAbsolutePath(),paramConfig);

        ArrayList<TextBlock> textBlocks = ocrResult.getTextBlocks();
        // 后处理
        List<String> strings = postProcessOcrResults(textBlocks);

        //清除临时文件
        if(file.exists()){
            file.delete();
        }

        return String.join("\n", strings);

    }

    /**
     * 后处理
     * @param textBlocks 文字块列表
     * @return 处理后的文字列表
     */
    public static List<String> postProcessOcrResults(List<TextBlock> textBlocks) {
        //重写排序规则
        textBlocks.sort((o1, o2) -> {

            //获取中心点的二倍坐标
            int x1 = o1.getBoxPoint().get(0).getX() + o1.getBoxPoint().get(2).getX();
            int y1 = o1.getBoxPoint().get(0).getY() + o1.getBoxPoint().get(2).getY();
            int x2 = o2.getBoxPoint().get(0).getX() + o2.getBoxPoint().get(2).getX();
            int y2 = o2.getBoxPoint().get(0).getY() + o2.getBoxPoint().get(2).getY();

            //获取两个文字块的高之和
            int temp1 = o2.getBoxPoint().get(3).getY() - o2.getBoxPoint().get(0).getY();
            int temp2 = o1.getBoxPoint().get(3).getY() - o1.getBoxPoint().get(0).getY();
            int temp = temp1 + temp2;

            //若中心点相差小于文字高度,视为统一行,按照列坐标排序
            y1 = y1 / temp;
            y2 = y2 / temp;
            if(y1 == y2){
                return x1 - x2;
            }
            return y1 - y2;
        });

        return textBlocks.stream()
                .map(TextBlock::getText)
                .collect(Collectors.toList());
    }



}
